{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in C:\\Users\\Administrator/.cache\\torch\\hub\\pytorch_vision_v0.4.2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1000])\n",
      "tensor(244)\n"
     ]
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"Untitled16.ipynb\n",
    "\n",
    "Automatically generated by Colaboratory.\n",
    "\n",
    "Original file is located at\n",
    "    https://colab.research.google.com/drive/1PBNv8qCKce_i_BG_kFjDUSuiEqNtcY1U\n",
    "\"\"\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.hub import load_state_dict_from_url\n",
    "\n",
    "\n",
    "model_urls = {\n",
    "    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n",
    "    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n",
    "    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n",
    "    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n",
    "    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n",
    "    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',\n",
    "    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',\n",
    "    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',\n",
    "    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',\n",
    "}\n",
    "\n",
    "def conv3x3(in_planes, out_planes, stride=1, padding=1):\n",
    "  return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=padding, bias=False)  #? Why no bias\n",
    "\n",
    "def conv1x1(in_planes, out_planes, stride=1):\n",
    "  return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) #? Why no bias: 如果卷积层之后是BN层，那么可以不用偏置参数，可以节省内存\n",
    "\n",
    "class BasicBlock(nn.Module):\n",
    "  expansion = 1 # 经过Block之后channel的变化量\n",
    "\n",
    "  def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):\n",
    "    # downsample: 调整维度一致之后才能相加\n",
    "    # norm_layer：batch normalization layer\n",
    "    super(BasicBlock, self).__init__()\n",
    "    if norm_layer is None:\n",
    "      norm_layer = nn.BatchNorm2d # 如果bn层没有自定义，就使用标准的bn层\n",
    "    self.conv1 = conv3x3(inplanes, planes, stride)\n",
    "    self.bn1 = norm_layer(planes)\n",
    "    self.relu = nn.ReLU(inplace=True)\n",
    "    self.conv2 = conv3x3(planes, planes)\n",
    "    self.bn2 = norm_layer(planes)\n",
    "    self.downsample = downsample\n",
    "    self.stride = stride\n",
    "\n",
    "  def forward(self, x):\n",
    "    identity = x  # 保存x\n",
    "\n",
    "    out = self.conv1(x)\n",
    "    out = self.bn1(out)\n",
    "    out = self.relu(out)\n",
    "\n",
    "    out = self.conv2(out)\n",
    "    out = self.bn2(out)\n",
    "\n",
    "    if self.downsample is not None:\n",
    "      identity = self.downsample(x)  # downsample调整x的维度，F(x)+x一致才能相加\n",
    "    \n",
    "    out += identity\n",
    "    out = self.relu(out) # 先相加再激活\n",
    "\n",
    "    return out\n",
    "\n",
    "class Bottleneck(nn.Module):\n",
    "  expansion = 4\n",
    "\n",
    "  def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):\n",
    "    super(Bottleneck, self).__init__()\n",
    "    if norm_layer is None:\n",
    "      norm_layer = nn.BatchNorm2d\n",
    "    \n",
    "    self.conv1 = conv1x1(inplanes, planes)\n",
    "    self.bn1 = norm_layer(planes)\n",
    "    self.conv2 = conv3x3(planes, planes, stride)\n",
    "    self.bn2 = norm_layer(planes)\n",
    "    self.conv3 = conv1x1(planes, planes * self.expansion) # 输入的channel数：planes * self.expansion\n",
    "    self.bn3 = norm_layer(planes * self.expansion)\n",
    "    self.relu = nn.ReLU(inplace=True)\n",
    "    self.downsample = downsample\n",
    "    self.stride = stride\n",
    "\n",
    "  def forward(self, x):\n",
    "    identity = x\n",
    "\n",
    "    out = self.conv1(x)\n",
    "    out = self.bn1(out)\n",
    "    out = self.relu(out)\n",
    "\n",
    "    out = self.conv2(out)\n",
    "    out = self.bn2(out)\n",
    "    out = self.relu(out)\n",
    "\n",
    "    out = self.conv3(out)\n",
    "    out = self.bn3(out)\n",
    "    \n",
    "    if self.downsample is not None:\n",
    "      identity = self.downsample(x)\n",
    "\n",
    "    out += identity\n",
    "    out = self.relu(out)\n",
    "\n",
    "    return out\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "  def __init__(self, block, layers, num_class=1000, norm_layer=None):\n",
    "    super(ResNet, self).__init__()\n",
    "    if norm_layer is None:\n",
    "      norm_layer = nn.BatchNorm2d\n",
    "    self._norm_layer = norm_layer\n",
    "\n",
    "    self.inplanes = 64\n",
    "\n",
    "    # conv1 in ppt figure\n",
    "    self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "    self.bn1 = norm_layer(self.inplanes)\n",
    "    self.relu = nn.ReLU(inplace=True)\n",
    "    self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "    self.layer1 = self._make_layer(block, 64, layers[0])\n",
    "    self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n",
    "    self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n",
    "    self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n",
    "    self.avgpool = nn.AdaptiveAvgPool2d((1,1))  # (1,1)等于GAP\n",
    "    self.fc = nn.Linear(512*block.expansion, num_class)\n",
    "\n",
    "    for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "  def _make_layer(self, block, planes, blocks, stride=1):\n",
    "    # 生成不同的stage/layer\n",
    "    # block: block type(basic block/bottle block)\n",
    "    # blocks: blocks的数量\n",
    "    norm_layer = self._norm_layer\n",
    "    downsample = None\n",
    "\n",
    "    if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "      # 需要调整维度\n",
    "      downsample = nn.Sequential(\n",
    "          conv1x1(self.inplanes, planes * block.expansion, stride),  # 同时调整spatial(H x W))和channel两个方向\n",
    "          norm_layer(planes * block.expansion)\n",
    "      )\n",
    "\n",
    "    layers = []\n",
    "    layers.append(block(self.inplanes, planes, stride, downsample, norm_layer)) # 第一个block单独处理\n",
    "    self.inplanes = planes * block.expansion  # 记录layerN的channel变化，具体请看ppt resnet表格\n",
    "    for _ in range(1, blocks): # 从1开始循环，因为第一个模块前面已经单独处理\n",
    "      layers.append(block(self.inplanes, planes, norm_layer=norm_layer))\n",
    "    return nn.Sequential(*layers)  # 使用Sequential层组合blocks，形成stage。如果layers=[2,3,4]，那么*layers=？\n",
    "\n",
    "  def forward(self, x):\n",
    "    x = self.conv1(x)\n",
    "    x = self.bn1(x)\n",
    "    x = self.relu(x)\n",
    "    x = self.maxpool(x)\n",
    "\n",
    "    x = self.layer1(x)\n",
    "    x = self.layer2(x)\n",
    "    x = self.layer3(x)\n",
    "    x = self.layer4(x)\n",
    "\n",
    "    x = self.avgpool(x)\n",
    "    x = torch.flatten(x, 1)\n",
    "    x = self.fc(x)\n",
    "\n",
    "    return x\n",
    "\n",
    "def _resnet(arch, block, layers, pretrained, progress, **kwargs):\n",
    "    model = ResNet(block, layers, **kwargs)\n",
    "    if pretrained:\n",
    "        state_dict = load_state_dict_from_url(model_urls[arch],\n",
    "                                              progress=progress)\n",
    "        model.load_state_dict(state_dict)\n",
    "    return model\n",
    "\n",
    "def resnet18(pretrained=False, progress=True, **kwargs):\n",
    "    r\"\"\"ResNet-18 model from\n",
    "    `\"Deep Residual Learning for Image Recognition\" <https://arxiv.org/pdf/1512.03385.pdf>`_\n",
    "    Args:\n",
    "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
    "        progress (bool): If True, displays a progress bar of the download to stderr\n",
    "    \"\"\"\n",
    "    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, #resnet单独写参数18的2 2 2  2 \n",
    "                   **kwargs)\n",
    "\n",
    "def resnet50(pretrained=False, progress=True, **kwargs):  #50层以上使用Bottleneck,50层一下使用BasicBlock\n",
    "    r\"\"\"ResNet-50 model from\n",
    "    `\"Deep Residual Learning for Image Recognition\" <https://arxiv.org/pdf/1512.03385.pdf>`_\n",
    "    Args:\n",
    "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
    "        progress (bool): If True, displays a progress bar of the download to stderr\n",
    "    \"\"\"\n",
    "    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,\n",
    "                   **kwargs)\n",
    "\n",
    "def resnet152(pretrained=False, progress=True, **kwargs):\n",
    "    r\"\"\"ResNet-50 model from\n",
    "    `\"Deep Residual Learning for Image Recognition\" <https://arxiv.org/pdf/1512.03385.pdf>`_\n",
    "    Args:\n",
    "        pretrained (bool): If True, returns a model pre-trained on ImageNet\n",
    "        progress (bool): If True, displays a progress bar of the download to stderr\n",
    "    \"\"\"\n",
    "    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,\n",
    "                   **kwargs)\n",
    "\n",
    "# model = resnet18(pretrained=True)\n",
    "# model.eval()\n",
    "\n",
    "import torch\n",
    "# model2 = torch.hub.load('pytorch/vision:v0.4.2', 'resnet18', pretrained=True)\n",
    "# or any of these variants\n",
    "# model = torch.hub.load('pytorch/vision:v0.4.2', 'resnet34', pretrained=True)\n",
    "# model = torch.hub.load('pytorch/vision:v0.4.2', 'resnet50', pretrained=True)\n",
    "# model = torch.hub.load('pytorch/vision:v0.4.2', 'resnet101', pretrained=True)\n",
    "model2 = torch.hub.load('pytorch/vision:v0.4.2', 'resnet152', pretrained=True)\n",
    "model2.eval()\n",
    "\n",
    "# model.state_dict()\n",
    "\n",
    "model2.state_dict()\n",
    "\n",
    "# Download an example image from the pytorch website\n",
    "import urllib\n",
    "# url, filename = (\"https://github.com/pytorch/hub/raw/master/dog.jpg\", \"dog.jpg\")\n",
    "# try: urllib.URLopener().retrieve(url, filename)\n",
    "# except: urllib.request.urlretrieve(url, filename)\n",
    "\n",
    "# sample execution (requires torchvision)\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "input_image = Image.open(\"dog.jpg\")\n",
    "preprocess = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])\n",
    "input_tensor = preprocess(input_image)\n",
    "input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model\n",
    "\n",
    "# move the input and model to GPU for speed if available\n",
    "if torch.cuda.is_available():\n",
    "    input_batch = input_batch.to('cuda')\n",
    "    model.to('cuda')\n",
    "\n",
    "with torch.no_grad():\n",
    "    output = model2(input_batch)\n",
    "# Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes\n",
    "# print(output[0])\n",
    "# The output has unnormalized scores. To get probabilities, you can run a softmax on it.\n",
    "\n",
    "result = torch.nn.functional.softmax(output[0], dim=0)\n",
    "\n",
    "result.argmax()\n",
    "print(result.shape)\n",
    "print(result.argmax())\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
